import data_handler
import runner
import asr_model_module
import torch.nn as nn
import torch
import config
from torch import optim

if __name__ == '__main__':
    loss_fun = data_handler.get_loss()
    train_iter = data_handler.get_iter()
    tokenizer = data_handler.PingyinVocabTokenizer()
    tokenizer.build_vocab()
    model = asr_model_module.AsrModel(fbank_dim=100, vocab_size=len(tokenizer), hidden_dim=1024, dropout=0.0897)
    lr = 1e-4
    optimizer = optim.Adam(model.parameters(), lr=lr)
    lr_scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=5, T_mult=2)
    runner_man = runner.RunnerGxl(model, optimizer, loss_fun, train_iter, logger=config.get_logger(), multi=False,
                                  device=torch.device('cpu'))
    runner_man.run(10)
